# this file contains some auxiliary code for results that are not in the paper or appendix but have helped
# me to think about the structural results.

using Statistics
using DataFrames

include("cenf.jl")

# Z^1

f= load("../output/ms_gm_output_withfixtheta_refined.jld2")

ms_fvals = f["ms_fvals"];
ms_θ = f["ms_θ"];
ms_β = f["ms_β"];
ms_η = f["ms_η"];
ms_γ = f["ms_γ"];
ms_logT = f["ms_logT"];
ms_logS = f["ms_logS"];

idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));

bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);

(changeArray_FGM,ioshares_before_FGM,ioshares_after_FGM) = cenf.solveCountries(bestθ, bestβ, bestη, bestγ, 1.0, exp.(bestlogT), exp.(bestlogS),"","")

println("Average: $(mean(changeArray_FGM[:,1]))")

# Z^2

f=load("../output/ms_f_output_withfixtheta_refined.jld2");
ms_fvals = f["ms_fvals"];
ms_θ = f["ms_θ"];
ms_β = f["ms_β"];
ms_η = f["ms_η"];
ms_γ = f["ms_γ"];
ms_logT = f["ms_logT"];
ms_logS = f["ms_logS"];

idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));

bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);

(changeArray_FF,ioshares_before_FF,ioshares_after_FF) = cenf.solveCountries(bestθ, bestβ, bestη, bestγ, 2.0, exp.(bestlogT), exp.(bestlogS),"", "")
println("Average: $(mean(changeArray_FF[:,1]))")

δarray = cenf.δ[:,1,1]
dfPlot = DataFrame(delta = δarray,dU = changeArray_FF[:,1], Legend = "")
p = Gadfly.plot(dfPlot, y=changeArray_FF[:,1], x=vec(δarray) , color="Legend", Geom.point, shape="Legend", Guide.title(""),
    Guide.xlabel("Enforcement cost"),
    Guide.ylabel("Counterfactual welfare increase, in percent"),
    Theme(default_color = colorant"red",
           highlight_width = 0pt)
     )

# STUDY THE OUTPUT

# 1.) plot the gamma_ni against the averages of X_ni/X_n (Referee 2)

cenf.lhs # c, n, i
avgXniXn = zeros(Float64, 35,35)
for n = 1:35
    for i = 1:35
        avgXniXn[n,i] = mean(cenf.lhs[:,n,i])
    end
end
bestγ
avgXniXn
p = Gadfly.plot(x = vec(log.(bestγ)) , y = vec(log.(avgXniXn)), Geom.point,
    Guide.xlabel("log estimated gamma_ni"),
    Guide.ylabel("log mean X_ni/X_n"),
    Theme(point_size = 2pt, highlight_width = 0pt))
draw(SVG("../output/comparison_gamma.svg",15cm, 10cm),p)
run(`cairosvg ../output/comparison_gamma.svg -o ../output/comparison_gamma.pdf`)

# 2.) compare z_gm and z_f
p = Gadfly.plot(y = vec(log.(cenf.z_f[1,:,:])) , x = vec(log.(cenf.z_gm[1,:,:])), Geom.point,
Guide.xlabel("log z^(1)"),
Guide.ylabel("log z^(2)"),
Theme(point_size = 2pt, highlight_width = 0pt))
draw(SVG("../output/comparison_z.svg",15cm, 10cm),p)
run(`cairosvg ../output/comparison_z.svg -o ../output/comparison_z.pdf`)

# 3.) study the shape of the d^I function

f= load("../output/ms_gm_output_withfixtheta_refined.jld2")
ms_fvals = f["ms_fvals"];
ms_θ = f["ms_θ"];
ms_β = f["ms_β"];
ms_η = f["ms_η"];
ms_γ = f["ms_γ"];
ms_logT = f["ms_logT"];
ms_logS = f["ms_logS"];
idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));
bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);
bestη_gm = deepcopy(bestη)
z = cenf.z_gm[1,:,:]
dI_gm = vec(1.0 .+ bestη_gm[1].*z.^bestη_gm[2])
p = Gadfly.plot(y = dI_gm , x = vec(z), Geom.point,
    Guide.xlabel("z_gm"),
    Guide.ylabel("dI"),
    Theme(point_size = 2pt, highlight_width = 0pt))

δ = cenf.δ[:,1,1]
dF = 1.0 ./ (1.0 .- bestβ[1].*δ .- bestβ[2].*δ.*δ)

f=load("../output/ms_f_output_withfixtheta_refined.jld2");
ms_fvals = f["ms_fvals"];
ms_θ = f["ms_θ"];
ms_β = f["ms_β"];
ms_η = f["ms_η"];
ms_γ = f["ms_γ"];
ms_logT = f["ms_logT"];
ms_logS = f["ms_logS"];
idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));
bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);
bestη_f = deepcopy(bestη)
z = cenf.z_f[1,:,:]
dI_f = vec(1.0 .+ bestη_f[1].*z.^bestη_f[2] )
p = Gadfly.plot(y = dI_f , x = vec(z), Geom.point,
    Guide.xlabel("z_f"),
    Guide.ylabel("dI"),
    Theme(point_size = 2pt, highlight_width = 0pt))
δ = cenf.δ[:,1,1]
dF = 1.0 ./ (1.0 .- bestβ[1].*δ .- bestβ[2].*δ.*δ)
p = Gadfly.plot(y = vec(dF) , x = vec(δ), Geom.point,
    Guide.xlabel("delta"),
    Guide.ylabel("dF"),
    Theme(point_size = 2pt, highlight_width = 0pt))

# compare
df1 = DataFrame(z = vec(cenf.z_gm[1,:,:]), dI = dI_gm, label="z_gm")
df2 = DataFrame(z = vec(cenf.z_f[1,:,:]), dI = dI_f, label="z_f")
dfplot = vcat(df1,df2)
p = Gadfly.plot(dfplot, x=:z, y=:dI , color="label", Geom.point, shape="label",
    Guide.xlabel("z"),
    Guide.ylabel("dI"),
    Theme(default_color = colorant"red",
           highlight_width = 0pt)
     )

p = Gadfly.plot(y = vec(dI_f) , x = vec(dI_gm), Geom.point,
Guide.xlabel("dI_gm"),
Guide.ylabel("dI_f"),
Theme(point_size = 2pt, highlight_width = 0pt))

# Experimentation time! **************************************

# 1.) Comparison of GM results with dI_GM and GM results with dI_F shows that the culprit is really dI: the latter is similar
#   to the F results with dI_F
# 2.) It's not the large values of dI that are driving the result. Capping both dI_GM and dI_F at 2.0 changes virtually nothing.

include("cenf.jl")
cenf.readData("C:/Users/johannes.boehm/Dropbox/cenf/data/revision-2017")

# Z^1

f= load("../output/ms_gm_output_withfixtheta_refined.jld2")

ms_fvals = f["ms_fvals"];
ms_θ = f["ms_θ"];
ms_β = f["ms_β"];
ms_η = f["ms_η"];
ms_γ = f["ms_γ"];
ms_logT = f["ms_logT"];
ms_logS = f["ms_logS"];

idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));

bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη_gm = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);

# BASELINE
dI_gm = 1.0 .+ bestη_gm[1].* cenf.z_gm.^bestη_gm[2]
(changeArray_FGM,ioshares_before_FGM,ioshares_after_FGM) = cenf.solveCountries(bestθ, bestβ, dI_gm, bestγ, exp.(bestlogT), exp.(bestlogS),"","")
# plot it
deltavec = cenf.δ[:,1,1]
xticks = [0.25, 0.5, 0.75, 1]
df = DataFrame(x = deltavec, y = 100.0 .* changeArray_FGM[:,1], name = vec(cenf.countrynames))
p=Gadfly.plot(df, x=:x, y=:y, Geom.point, label =:name, Geom.label, 
    Coord.Cartesian(xmin=0.08,xmax=1.08), Guide.xticks(ticks=xticks), 
    Guide.xlabel("Enforcement cost"), Guide.ylabel("Counterfactual welfare increase, in percent",orientation=:vertical), 
    Guide.title(""))
# capping dI at
dI = 1.0 .+ bestη_gm[1].* cenf.z_gm.^bestη_gm[2]
for c=1:109
    for n=1:35
        for i=1:35
            dI[c,n,i] = (dI[c,n,i] > 2.0) ? 2.0 : dI[c,n,i]
        end
    end
end
(changeArray_FGM,ioshares_before_FGM,ioshares_after_FGM) = cenf.solveCountries(bestθ, bestβ, dI, bestγ, exp.(bestlogT), exp.(bestlogS),"","")
# plot it
deltavec = cenf.δ[:,1,1]
xticks = [0.25, 0.5, 0.75, 1]
df = DataFrame(x = deltavec, y = 100.0 .* changeArray_FGM[:,1], name = vec(cenf.countrynames))
p=Gadfly.plot(df, x=:x, y=:y, Geom.point, label =:name, Geom.label, 
    Coord.Cartesian(xmin=0.08,xmax=1.08), Guide.xticks(ticks=xticks), 
    Guide.xlabel("Enforcement cost"), Guide.ylabel("Counterfactual welfare increase, in percent",orientation=:vertical), 
    Guide.title(""))

println("Average: $(mean(changeArray_FGM[:,1]))")

# Z^2

f=load("../output/ms_f_output_withfixtheta_refined.jld2");
ms_fvals = f["ms_fvals"];
ms_θ = f["ms_θ"];
ms_β = f["ms_β"];
ms_η = f["ms_η"];
ms_γ = f["ms_γ"];
ms_logT = f["ms_logT"];
ms_logS = f["ms_logS"];

idx_best = findfirst(isequal(maximum(ms_fvals)), vec(ms_fvals));

bestθ = ms_θ[idx_best]
bestβ = vec(ms_β[idx_best,:])
bestη_f = vec(ms_η[idx_best,:])
bestγ = reshape(ms_γ[idx_best,:],35,35);
bestlogT = reshape(ms_logT[idx_best,:],109,35);
bestlogS = reshape(ms_logS[idx_best,:],109,35);

dI_f = 1.0 .+ bestη_f[1].* cenf.z_f .+ bestη_f[2].* cenf.z_f.* cenf.z_f
(changeArray_FGM,ioshares_before_FGM,ioshares_after_FGM) = cenf.solveCountries(bestθ, bestβ, dI_f, bestγ, exp.(bestlogT), exp.(bestlogS),"","")

dI_f = 1.0 .+ bestη_f[1].* cenf.z_f.^bestη_f[2]
for c=1:109
    for n=1:35
        for i=1:35
            dI_f[c,n,i] = (dI_f[c,n,i] > 2.0) ? 2.0 : dI_f[c,n,i]
        end
    end
end
(changeArray_FF,ioshares_before_FF,ioshares_after_FF) = cenf.solveCountries(bestθ, bestβ, dI_f, bestγ, exp.(bestlogT), exp.(bestlogS),"","")
# plot it
deltavec = cenf.δ[:,1,1]
xticks = [0.25, 0.5, 0.75, 1]
df = DataFrame(x = deltavec, y = 100.0 .* changeArray_FF[:,1], name = vec(cenf.countrynames))
p=Gadfly.plot(df, x=:x, y=:y, Geom.point, label =:name, Geom.label, 
    Coord.Cartesian(xmin=0.08,xmax=1.08), Guide.xticks(ticks=xticks), 
    Guide.xlabel("Enforcement cost"), Guide.ylabel("Counterfactual welfare increase, in percent",orientation=:vertical), 
    Guide.title(""))

# compare
df1 = DataFrame(z = vec(cenf.z_gm[1,:,:]), dI = dI_gm, label="z_gm")
df2 = DataFrame(z = vec(cenf.z_f[1,:,:]), dI = vec(dI_f[1,:,:]), label="z_f")
dfplot = vcat(df1,df2)
p = Gadfly.plot(df2, x=:z, y=:dI , color="label", Geom.point, shape="label",
    Guide.xlabel("z"),
    Guide.ylabel("dI"),
    Theme(default_color = colorant"red",
           highlight_width = 0pt)
     )
p = Gadfly.plot(x=vec(dI_gm), y=vec(dI_f), Geom.point,
Guide.xlabel("dI_gm"),
Guide.ylabel("dI_f")
)
tempdf = DataFrame(dI_gm = vec(dI_gm[1,:,:]), dI_f = vec(dI_f[1,:,:]))
p = Gadfly.plot(tempdf[(tempdf[:dI_gm] .> 1.15) .& (tempdf[:dI_f] .> 1.15),:], x=:dI_gm, y=:dI_f,  Geom.point,
Guide.xlabel("dI_gm"),
Guide.ylabel("dI_f"),
intercept = [0.0], slope = [1.0], Geom.abline(style=[:solid])
)

using FixedEffectModels
df = DataFrame(ratio = vec(log.((cenf.z_gm[1,:,:] .- 0.001) ./ (cenf.z_f[1,:,:] .- 0.001))), XniXn = vec(log.(cenf.lhs[106,:,:])))
filter!(r -> !isnan(r.ratio), df)
m = reg(df, @formula(ratio ~ XniXn))
m.coef[1]
m.coef[2]

p = Gadfly.plot(df, x=:ratio, y=:XniXn,  Geom.point,
Guide.ylabel("log Xni/Xn"),
Guide.xlabel("log z^1/z^2")
)
draw(SVG("../output/comparison_XniXn.svg",15cm, 10cm),p)
run(`cairosvg ../output/comparison_XniXn.svg -o ../output/comparison_XniXn.pdf`)
